#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import rospy, sys
import moveit_commander
from geometry_msgs.msg import PoseStamped

class MoveItIkDemo:
    def __init__(self):
        moveit_commander.roscpp_initialize(sys.argv)
        rospy.init_node('moveit_ik_demo')

        arm = moveit_commander.MoveGroupCommander('robot_arm')
        end_effector_link = arm.get_end_effector_link()
        reference_frame = 'base_link'
        arm.set_pose_reference_frame(reference_frame)
        arm.allow_replanning(True)
        arm.set_goal_position_tolerance(0.001)
        arm.set_goal_orientation_tolerance(0.01)
        arm.set_max_acceleration_scaling_factor(0.1)
        arm.set_max_velocity_scaling_factor(0.1)

        # 回到 home 位
        # arm.set_named_target('home')
        # arm.go(wait=True)
        # rospy.sleep(2)

        # 设置目标位姿
        target_pose = PoseStamped()
        target_pose.header.frame_id = reference_frame
        target_pose.header.stamp = rospy.Time.now()
        target_pose.pose.position.x = -0.1960999951310097
        target_pose.pose.position.y = 0.17977465457891884
        target_pose.pose.position.z = 0.10555906712371071
        target_pose.pose.orientation.x = 0.70827
        target_pose.pose.orientation.y = -0.70578
        target_pose.pose.orientation.z = 0.0093297
        target_pose.pose.orientation.w = -0.011898

        # 持续规划直到用户确认执行
        while not rospy.is_shutdown():
            rospy.loginfo("正在进行运动规划...")
            arm.set_start_state_to_current_state()
            arm.set_pose_target(target_pose, end_effector_link)
            plan_success, traj, planning_time, error_code = arm.plan()

            if plan_success:
                rospy.loginfo("✅ 规划成功！请在 RViz 中查看轨迹。")
                user_input = input("是否执行该轨迹？(y=执行 / n=重新规划 / q=退出): ").strip().lower()

                if user_input == 'y':
                    rospy.loginfo("开始执行规划轨迹...")
                    arm.execute(traj, wait=True)
                    rospy.loginfo("执行完成。")
                    break
                elif user_input == 'n':
                    rospy.logwarn("重新规划中...")
                    continue
                elif user_input == 'q':
                    rospy.logwarn("用户选择退出。")
                    break
                else:
                    rospy.logwarn("无效输入，请重新输入。")
                    continue
            else:
                rospy.logerr("❌ 规划失败，正在重新规划...")
                rospy.sleep(1.0)

        moveit_commander.roscpp_shutdown()
        moveit_commander.os._exit(0)

if __name__ == "__main__":
    MoveItIkDemo()
